from gensim import models, similarities
from model.SentenceSimilarityModel import SentenceSimilarityModel


class ModelTrainHandler(object):

    def __init__(self, daraPreHandler):
        self.daraPreHandler = daraPreHandler

    # tfidf模型
    def TfidfModel(self):
        self.daraPreHandler.preModel()
        simpleCorpus = self.daraPreHandler.simpleCorpus

        # 转换模型
        model = models.TfidfModel(simpleCorpus)
        corpus = model[simpleCorpus]

        # 创建相似度矩阵
        index = similarities.MatrixSimilarity(corpus)

        dictionary = self.daraPreHandler.dictionary
        sentences = self.daraPreHandler.sentences

        return SentenceSimilarityModel(model, dictionary, corpus, index, sentences)

    # lsi模型
    def LsiModel(self):
        self.daraPreHandler.preModel()
        simpleCorpus = self.daraPreHandler.simpleCorpus

        # 转换模型
        model = models.LsiModel(simpleCorpus)
        corpus = model[simpleCorpus]

        # 创建相似度矩阵
        index = similarities.MatrixSimilarity(corpus)

        dictionary = self.daraPreHandler.dictionary
        sentences = self.daraPreHandler.sentences

        return SentenceSimilarityModel(model, dictionary, corpus, index, sentences)

    #lda模型
    def LdaModel(self):
        self.daraPreHandler.preModel()
        simpleCorpus = self.daraPreHandler.simpleCorpus

        # 转换模型
        model = models.LdaModel(simpleCorpus)
        corpus = model[simpleCorpus]

        # 创建相似度矩阵
        index = similarities.MatrixSimilarity(corpus)

        dictionary = self.daraPreHandler.dictionary
        sentences = self.daraPreHandler.sentences

        return SentenceSimilarityModel(model, dictionary, corpus, index, sentences)
